from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModel
import gradio as gr
import torch
import transformers

print(torch.__version__)
print(torch.cuda.is_available())
print(transformers.__version__)

# 清除gpu
torch.cuda.empty_cache()

path = r'M:\moudels\chatglm3\chatglm3-6b'

tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True, use_fast=True)

model = AutoModel.from_pretrained(path, trust_remote_code=True, device_map="auto", load_in_4bit=True,
                                  bnb_4bit_compute_dtype=torch.float16)
model = model.eval()
CSS = """
*{
    font-family: bold STKaiti, Kaiti SC, Kaiti, BlinkMacSystemFont, Helvetica Neue, PingFang SC, Microsoft YaHei, Source Han Sans SC, Noto Sans CJK SC, WenQuanYi Micro Hei, Arial, sans-serif;
    font-size: 20px;
    color: rgb(147 102 0);
}
"""

prompt = """
餐厅客服机器人
1 明确机器人的角色
--你是一个订餐机器人，为面馆餐厅自动收集订餐信息
2 明确机器人跟客户交流的流程
--你要首先问候顾客，然后跟用户沟通，并通过用户的回复收集订单信息，
--收集完成后需要确认顾客是否需要添加其他内容，
--询问是否自取或者外送， 需要外送的，你需要询问地址，不需要外送则到最后的步骤。
--完成上面流程最后的步骤你需要告诉顾客订单总金额，并送上祝福
3 体统流程具体相关信息
-- 请确保明确所有的面的名称，规格，价格信息
-- 面类
---- 名称：油泼面 规格：大 价钱：13
---- 名称：油泼面 规格：中 价钱：12
---- 名称：油泼面 规格：小 价钱：11
---- 名称：肉臊子干拌 规格：大 价钱：16
---- 名称：肉臊子干拌 规格：中 价钱：15
---- 名称：肉臊子干拌 规格：小 价钱：14
---- 名称：素臊子干拌 规格：大 价钱：22
---- 名称：素臊子干拌 规格：中 价钱：21
---- 名称：素臊子干拌 规格：小 价钱：20
-- 饮料
---- 名称：大窑 价钱：5
---- 名称：冰峰 价钱：3
---- 名称：可乐 价钱：3.5
"""



dome = gr.Blocks(css=CSS)


with dome:
    # 2. 流式对话
    def chat_stream_gsy(mes, current_his):
        print("流式对话模式:")
        response = ""
        global prompt
        # history = current_his
        # current_his.append()
        for response, history in model.stream_chat(tokenizer,
                                                   mes,
                                                   history=current_his[-1:],
                                                   prompt=[prompt],
                                                   temperature=0.8,  # 控制随机性 (0-1)
                                                   top_p=0.9,  # 核采样参数
                                                   max_length=2048,  # 最大生成长度
                                                   repetition_penalty=1.2  # 重复惩罚系数
                                                   ):
            # 实时打印新增内容
            yield response
            # print(new_text, end="", flush=True)
        print("response:  ", response)
        print("current_his:  ", current_his)

    chat_interface = gr.ChatInterface(
        fn=chat_stream_gsy,
        type="messages",
        title="餐厅客服机器人",
        description="与ChatGLM3-6B模型进行实时对话",
        submit_btn="提交",
        stop_btn="撤销",
    )

if __name__ == '__main__':
    pass
    dome.launch()
    # chat_stream_gsy("你好", [])
